import cv2
import mediapipe as mp
# ========== 配置区域 ==========
# 文件参数
INPUT_PATH = "input.mp4"        # 输入视频路径
OUTPUT_PATH = "output.mp4"       # 输出视频路径

# 模型参数
MIN_DETECTION_CONFIDENCE = 0.5   # 最小检测置信度 (0.0-1.0)
MIN_TRACKING_CONFIDENCE = 0.5    # 最小跟踪置信度 (0.0-1.0)
MODEL_COMPLEXITY = 2             # 模型复杂度 (0-2)

# 绘图参数
COLOR_BGR = (0, 255, 0)         # 骨骼线颜色 (BGR格式)
THICKNESS = 3                   # 骨骼线粗细
DRAW_KEYPOINTS = False          # 是否绘制关键点
KEYPOINT_SIZE = 2               # 关键点半径（当DRAW_KEYPOINTS=True时生效）

# 视频编码参数
FOURCC = 'mp4v'                 # 视频编码器 (常见：'mp4v', 'avc1', 'XVID')
FORCE_COLOR_MODE = True         # 强制使用彩色模式输出
# =============================

def draw_skeleton_on_video():
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(
        min_detection_confidence=MIN_DETECTION_CONFIDENCE,
        min_tracking_confidence=MIN_TRACKING_CONFIDENCE,
        model_complexity=MODEL_COMPLEXITY
    )

    # 配置绘图样式
    circle_radius = KEYPOINT_SIZE if DRAW_KEYPOINTS else None
    connection_style = {
        connection: mp.solutions.drawing_utils.DrawingSpec(
            color=COLOR_BGR,
            thickness=THICKNESS,
            circle_radius=circle_radius
        ) for connection in mp_pose.POSE_CONNECTIONS
    }

    # 初始化视频流
    cap = cv2.VideoCapture(INPUT_PATH)
    if not cap.isOpened():
        print(f"错误：无法打开视频源 {INPUT_PATH}")
        return

    # 获取视频参数
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    # 配置视频编码器
    fourcc = cv2.VideoWriter_fourcc(*FOURCC)
    out = cv2.VideoWriter(OUTPUT_PATH, fourcc, fps, 
                        (frame_width, frame_height),
                        isColor=FORCE_COLOR_MODE)

    try:
        while cap.isOpened():
            success, frame = cap.read()
            if not success:
                break
                
            # 处理帧数据
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = pose.process(rgb_frame)
            
            # 绘制骨骼
            if results.pose_landmarks:
                mp.solutions.drawing_utils.draw_landmarks(
                    frame,
                    results.pose_landmarks,
                    mp_pose.POSE_CONNECTIONS,
                    connection_drawing_spec=connection_style
                )
            
            out.write(frame)
    finally:
        # 确保资源释放
        cap.release()
        out.release()

if __name__ == "__main__":
    draw_skeleton_on_video()